import numpy as np
from scipy.integrate import odeint
from pdb import set_trace
"""
Functions to calculate the transient deformations, displacements
and forces from a Hertzian impact between a sphere and a load
cell.
"""
# Global constants
g = 9.8             # gravity (m/s^2)

def Hertz_constant(E1,E2,v1,v2,R):
    """
    function to calculate the effective stiffness used in the Hertz contact
    theory.
    E1 - tablet modulus (Pa)
    E2 - plane modulus (Pa)
    v1 - tablet Poisson ratio
    v2 - Plane Poisson ratio
    R - radius of sphere (m)
    return K - constant for Hertz force F = Ky^1.5
    """
    E_eff = 1. / ( (1.-v1**2)/E1 + (1.-v2**2)/E2 )
    K = 4./3.*R**0.5*E_eff
    return K  

def Hertz_force(y,K):
    """
    function to calculate the force on the load cell
    y - indentation depth (m)
    K - Hertz constant (kg/(m^0.5*s^2))
    return F - total contact force between tablet and load cell (N)
    """
    if np.isscalar(y):
        y = np.max([y,0])
    else:
        y[y<0] = 0
    F = K*y**(1.5)
    return F
     
def dissipative_constant(m,K,COR):
    """
    function to calculate the constant for the dissipative force 
    (see Tsuji Powder Tech, 71 (1992) 239)
    m - tablet mass (kg)
    K - Hertz constant (kg/(m^0.5*s^2))
    COR - coefficient of restitution for the tablet
    return h - constant for the dissipative force F = h*y**0.25*y_prime
    """
    a = -1*(np.sqrt(5)*np.log(COR)) / np.sqrt(np.log(COR)**2 + np.pi**2)
    h = a*np.sqrt(m*K)
    return h

def dissipative_force(eta,y,y_prime):
    """
    function to calculate dissipative force (see Tsuji Powder Tech, 71
    (1992) 239)
    eta - scalar dissipative factor
    y - indentation depth (m)
    y_prime - velocity (rate of change of indentation depth) (m/s)
    return F - dissipation force
    """
    if np.isscalar(y):
        y = np.max([y,0])
    else:
        y[y<0] = 0
    F = eta * y**0.25 * y_prime
    return F

def derivatives(y,t,m,K,eta):
    """
    function to calculate the derivatives
    y - dependent variable of coupled ODEs - 
    [indentation depth, rate change indentation depth] - [m,m/s]
    t - discrete time array for numerical integration (s)
    K -Hertz_constant (kg/(m^0.5*s^2))
    eta - dissipative constant kg*m^(1/4)/s
    return [2x1] numpy array derivatives
    """
    try:
        y0_prime = y[1]
        y1_prime = -1./m * (Hertz_force(y[0],K) 
                   + dissipative_force(eta,y[0],y[1]) - m*9.8)
        derivative = np.array([y0_prime,y1_prime])
    except (SystemExit, KeyboardInterrupt):
        raise
    return derivative 
    
def simulate(t=np.linspace(0,10**-3,10**5), m=0.407e-3, R=4.76e-3, 
        h=4.7e-2, Es=3.0e9, Ep=3.0e9, vs=0.33, vp=0.33, COR=0.8, cal=0.001):
    """
    Calculate the deformation, deformation rate and force
    at time points given in t for passed parameters.
    t - array of numpy time points (s)
    h - drop height (m)
    R - sphere radius (m)
    Es - sphere modulus (Pa)
    Ep - plane modulus (Pa)
    vs - sphere Poisson's ratio
    vp - plane Poisson's ratio
    m - sphere mass (kg)
    COR - coefficient of restitution
    cal - load cell calibration (N/mV)
    """
    K = Hertz_constant(Es,Ep,vs,vp,R)
    eta = dissipative_constant(m,K,COR)
    vo = g*(2*h/g)**0.5
    y_init = [0,vo]
    y = odeint(derivatives,y_init,t,(m,K,eta))
    mask = y[:,0] > 0
    t = t[mask]
    y = y[mask]
    pts = np.sum(mask)
    step = max(1, pts / 1000)
    t = np.array([t[i] for i in range(0,pts,step)])
    y = np.array([y[i,:] for i in range(0,pts,step)])
    f = Hertz_force(y[:,0],K) 
    mask = ~np.isnan(f)
    dynamics = np.zeros((t.shape[0]), dtype=[('time','<f4'), 
        ('deformation','<f4'), ('deformation rate', '<f4'), ('force', '<f4'),
        ('voltage','<f4')])
    dynamics['time'] = t 
    dynamics['deformation'] = y[:,0] 
    dynamics['deformation rate'] = y[:,1] 
    dynamics['force'] = f 
    dynamics['voltage'] = 1./ cal * dynamics['force']
    return dynamics
